# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from dataclasses import dataclass, field
from typing import Literal, Optional

from transformers import Seq2SeqTrainingArguments
from transformers.utils.versions import require_version

from swift.plugin import LOSS_MAPPING
from swift.trainers import TrainerFactory
from swift.trainers.arguments import TrainArgumentsMixin
from swift.utils import (add_version_to_work_dir, get_device_count, get_logger, get_pai_tensorboard_dir, is_master,
                         is_mp, is_pai_training_job, is_swanlab_available)
from .base_args import BaseArguments, to_abspath
from .tuner_args import TunerArguments

logger = get_logger()


@dataclass
class Seq2SeqTrainingOverrideArguments(TrainArgumentsMixin, Seq2SeqTrainingArguments):
    """重写 `Seq2SeqTrainingArguments` 中的默认参数配置"""

    # 指定输出目录（可选），用于保存模型、日志等
    output_dir: Optional[str] = None

    # 学习率（可选），如果不设置会根据 train_type 自动指定
    learning_rate: Optional[float] = None

    # 评估策略：'steps' 或 'epoch'，默认从 save_strategy 继承
    eval_strategy: Optional[str] = None

    # 是否使用 fp16 混合精度训练
    fp16: Optional[bool] = None

    # 是否使用 bf16 混合精度训练
    bf16: Optional[bool] = None

    # 初始化输出目录
    def _init_output_dir(self):
        # 如果未指定输出目录，则使用默认路径：output/{模型后缀名}
        if self.output_dir is None:
            self.output_dir = f'output/{self.model_suffix}'
        # 将路径转换为绝对路径
        self.output_dir = to_abspath(self.output_dir)

    # 初始化评估策略
    def _init_eval_strategy(self):
        # 如果未指定评估策略，默认使用保存策略（save_strategy）
        if self.eval_strategy is None:
            self.eval_strategy = self.save_strategy

        # 如果评估策略为 'no'，禁用评估，则清除评估步数、设置数据拆分比例为0
        if self.eval_strategy == 'no':
            self.eval_steps = None
            self.split_dataset_ratio = 0.
            logger.info(f'Setting args.split_dataset_ratio: {self.split_dataset_ratio}')
        # 如果策略是 steps 且评估步数没有指定，则跟随保存步数
        elif self.eval_strategy == 'steps' and self.eval_steps is None:
            self.eval_steps = self.save_steps

        # 设置 transformers 中真正生效的 evaluation_strategy 属性
        self.evaluation_strategy = self.eval_strategy

    # 初始化用于保存最佳模型的指标（如 ROUGE-L 或 loss）
    def _init_metric_for_best_model(self):
        # 如果未设置度量指标，生成模型则使用 ROUGE-L，否则用 loss
        if self.metric_for_best_model is None:
            self.metric_for_best_model = 'rouge-l' if self.predict_with_generate else 'loss'

    # 数据类初始化完成后自动调用的钩子方法
    def __post_init__(self):
        # 初始化输出目录
        self._init_output_dir()

        # 初始化评估指标
        self._init_metric_for_best_model()

        # 如果未明确指定 greater_is_better，根据度量指标名称推断
        if self.greater_is_better is None and self.metric_for_best_model is not None:
            # 如果度量指标中不包含 'loss'，则代表越大越好（如 ROUGE、BLEU）
            self.greater_is_better = 'loss' not in self.metric_for_best_model

        # 如果学习率未指定，则根据训练类型选择默认值
        if self.learning_rate is None:
            if self.train_type == 'full':
                self.learning_rate = 1e-5
            else:
                self.learning_rate = 1e-4

        # 初始化评估策略
        self._init_eval_strategy()


@dataclass
class SwanlabArguments:

    # swanlab 的访问令牌，用于认证
    swanlab_token: Optional[str] = None
    # swanlab 中的项目名称
    swanlab_project: Optional[str] = None
    # swanlab 工作空间名称
    swanlab_workspace: Optional[str] = None
    # 本次实验的名称，如果不设置则默认为输出路径
    swanlab_exp_name: Optional[str] = None
    # 运行模式，可选为 cloud（云端）或 local（本地），默认为 cloud
    swanlab_mode: Literal['cloud', 'local'] = 'cloud'

    # 初始化 swanlab 的方法，在训练脚本中被调用
    def _init_swanlab(self):
        # 如果没有安装 swanlab，则报错提示用户安装
        if not is_swanlab_available():
            raise ValueError('You are using swanlab as `report_to`, please install swanlab by '
                             '`pip install swanlab`')

        # 如果没有设置实验名称，则使用输出目录作为实验名称
        if not self.swanlab_exp_name:
            self.swanlab_exp_name = self.output_dir

        # 从 transformers 框架中导入集成回调的注册表
        from transformers.integrations import INTEGRATION_TO_CALLBACK
        # 导入 swanlab 库
        import swanlab
        # 从 swanlab 的 transformers 集成中导入回调类
        from swanlab.integration.transformers import SwanLabCallback

        # 如果设置了 swanlab_token，则调用 login 方法进行身份认证
        if self.swanlab_token:
            swanlab.login(self.swanlab_token)

        # 将 swanlab 的回调类注册到 transformers 框架中，便于训练时调用
        INTEGRATION_TO_CALLBACK['swanlab'] = SwanLabCallback(
            project=self.swanlab_project,          # swanlab 项目名称
            workspace=self.swanlab_workspace,      # swanlab 工作空间
            experiment_name=self.swanlab_exp_name, # swanlab 实验名称
            config={'UPPERFRAME': '🐦‍⬛ms-swift'},  # 自定义配置，用于标识框架来源
            mode=self.swanlab_mode,                # 运行模式（cloud 或 local）
        )

@dataclass
class TrainArguments(SwanlabArguments, TunerArguments, Seq2SeqTrainingOverrideArguments, BaseArguments):
    """
    训练参数类，继承了多个参数基类：
    - SwanlabArguments: SwanLab 可视化相关配置
    - TunerArguments: 超参数微调相关配置
    - Seq2SeqTrainingOverrideArguments: Seq2Seq 训练相关配置
    - BaseArguments: 基础参数配置

    参数说明：
    - add_version: 是否在输出目录添加版本信息，默认 True
    - resume_only_model: 是否仅恢复模型（不加载优化器等状态），默认 False
    - loss_type: 使用的损失函数类型，可在插件中定义
    - packing: 是否启用数据集 packing（打包多个样本进同一序列）
    - lazy_tokenize: 是否启用惰性分词
    - max_new_tokens: 最大生成新 token 数，默认 64
    - temperature: 采样温度，默认 0（贪婪）
    - optimizer / metric: 优化器与评估指标类型，由插件提供
    """

    # 是否给输出目录添加版本后缀
    add_version: bool = True
    # 是否仅恢复模型（跳过优化器、scheduler 等）
    resume_only_model: bool = False
    # 是否创建 checkpoint 的软链接
    create_checkpoint_symlink: bool = False

    # 数据集处理相关
    packing: bool = False
    lazy_tokenize: Optional[bool] = None

    # 插件机制支持的参数
    loss_type: Optional[str] = field(default=None, metadata={'help': f'loss_func choices: {list(LOSS_MAPPING.keys())}'})
    optimizer: Optional[str] = None
    metric: Optional[str] = None

    # 附加配置
    max_new_tokens: int = 64
    temperature: float = 0.
    load_args: bool = False  # 是否从已有配置中加载参数

    # DeepSpeed ZeRO++ 优化的分区大小
    zero_hpz_partition_size: Optional[int] = None

    # ------------------ 初始化 lazy_tokenize 的逻辑 ------------------ #
    def _init_lazy_tokenize(self):
        # streaming 模式下不能启用 lazy_tokenize，会禁用并给出警告
        if self.streaming and self.lazy_tokenize:
            self.lazy_tokenize = False
            logger.warning('Streaming and lazy_tokenize are incompatible. '
                           f'Setting args.lazy_tokenize: {self.lazy_tokenize}.')
        # 如果未手动指定，自动判断是否启用 lazy_tokenize（用于多模态非流式情况）
        if self.lazy_tokenize is None:
            self.lazy_tokenize = self.model_meta.is_multimodal and not self.streaming
            logger.info(f'Setting args.lazy_tokenize: {self.lazy_tokenize}')

    # ------------------ dataclass 初始化钩子 ------------------ #
    def __post_init__(self) -> None:
        # 如果启用了 packing，但未使用 flash attention，给出警告
        if self.packing and self.attn_impl != 'flash_attn':
            logger.warning('The "packing" feature needs to be used in conjunction with "flash_attn". '
                           'Please specify `--attn_impl flash_attn`.')

        # 检查是否要从 checkpoint 恢复
        if self.resume_from_checkpoint:
            self.resume_from_checkpoint = to_abspath(self.resume_from_checkpoint, True)
            if self.resume_only_model:
                if self.train_type == 'full':
                    self.model = self.resume_from_checkpoint
                else:
                    self.adapters = [self.resume_from_checkpoint]

        # 手动调用各父类的 __post_init__
        BaseArguments.__post_init__(self)
        Seq2SeqTrainingOverrideArguments.__post_init__(self)
        TunerArguments.__post_init__(self)

        # 如果未显式指定 optimizer，根据 lora 或 galore 策略自动选择
        if self.optimizer is None:
            if self.lorap_lr_ratio:
                self.optimizer = 'lorap'
            elif self.use_galore:
                self.optimizer = 'galore'

        # 数据集不能为空，否则抛出异常
        if len(self.dataset) == 0:
            raise ValueError(f'self.dataset: {self.dataset}, Please input the training dataset.')

        # 兼容 PAI（阿里云平台）的逻辑
        self._handle_pai_compat()

        # 初始化 DeepSpeed 相关配置
        self._init_deepspeed()

        # 初始化训练设备
        self._init_device()

        # 初始化 lazy_tokenize 策略
        self._init_lazy_tokenize()

        # 初始化分布式训练 accelerator 配置
        if getattr(self, 'accelerator_config', None) is None:
            self.accelerator_config = {'dispatch_batches': False}

        # 获取 Trainer 使用的标准 TrainingArguments
        self.training_args = TrainerFactory.get_training_args(self)
        self.training_args.remove_unused_columns = False

        # 添加版本信息到输出目录
        self._add_version()

        # 初始化 swanlab 日志可视化工具
        if 'swanlab' in self.report_to:
            self._init_swanlab()

    # ------------------ 初始化 DeepSpeed 配置 ------------------ #
    def _init_deepspeed(self):
        if self.deepspeed:
            require_version('deepspeed')  # 检查是否已安装
            if is_mp():  # mp 模式与 DeepSpeed 不兼容
                raise ValueError('DeepSpeed is not compatible with `device_map`. '
                                 f'n_gpu: {get_device_count()}, '
                                 f'local_world_size: {self.local_world_size}.')

            # 预定义 DeepSpeed 配置模板路径
            ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'ds_config'))
            deepspeed_mapping = {
                name: f'{name}.json'
                for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload']
            }

            # 根据命令行输入匹配配置文件路径
            for ds_name, ds_config in deepspeed_mapping.items():
                if self.deepspeed == ds_name:
                    self.deepspeed = os.path.join(ds_config_folder, ds_config)
                    break

            # 解析 json 配置为 dict
            self.deepspeed = self.parse_to_dict(self.deepspeed)

            # 设置 ZeRO++ 分区配置
            if self.zero_hpz_partition_size is not None:
                assert 'zero_optimization' in self.deepspeed
                self.deepspeed['zero_optimization']['zero_hpz_partition_size'] = self.zero_hpz_partition_size
                logger.warn('If `zero_hpz_partition_size`(ZeRO++) causes grad_norm NaN, please'
                            ' try `--torch_dtype float16`')
            logger.info(f'Using deepspeed: {self.deepspeed}')

    # ------------------ 兼容 PAI 平台训练 ------------------ #
    def _handle_pai_compat(self) -> None:
        if not is_pai_training_job():
            return
        logger.info('Handle pai compat...')
        pai_tensorboard_dir = get_pai_tensorboard_dir()
        if self.logging_dir is None and pai_tensorboard_dir is not None:
            self.logging_dir = pai_tensorboard_dir
            logger.info(f'Setting args.logging_dir: {self.logging_dir}')
        self.add_version = False  # 在 PAI 平台禁用版本追加
        logger.info(f'Setting args.add_version: {self.add_version}')

    # ------------------ 添加版本信息到输出路径 ------------------ #
    def _add_version(self):
        """为输出路径添加版本号，并设置日志路径"""
        if self.add_version:
            self.output_dir = add_version_to_work_dir(self.output_dir)
            logger.info(f'output_dir: {self.output_dir}')

        # 设置默认的日志目录
        if self.logging_dir is None:
            self.logging_dir = f'{self.output_dir}/runs'

        # 统一转为绝对路径
        self.logging_dir = to_abspath(self.logging_dir)
        if is_master():
            os.makedirs(self.output_dir, exist_ok=True)  # 创建输出目录

        # 设置 run_name（如果未手动指定）
        if self.run_name is None:
            self.run_name = self.output_dir

        # 设置 training_args 的关键路径信息
        self.training_args.output_dir = self.output_dir
        self.training_args.run_name = self.run_name
        self.training_args.logging_dir = self.logging_dir
